using System;
using System.Runtime.InteropServices;
using UnmanagedClass;

public unsafe struct BlockBuffer : IDisposable
{
    [StructLayout(LayoutKind.Sequential)]
    private struct Block
    {
        public Block* prev;
        public Block* next;
        public byte* writePosition;
        public byte* blockEnd;
        public fixed byte data[1]; // flexible array member
    }

    private const int minAllocationSize = 1024;
    private readonly int oneBlockSize;

    private Block* headBlock;
    private Block* tailBlock;

    public int blockCount { get; private set; }

    public bool IsCreated => tailBlock != null;

    public int MaxDataSizeOfBlock => oneBlockSize - 4 * sizeof(byte*);
    public int AllDataSize => MaxDataSizeOfBlock * (blockCount - 1) + (int)(tailBlock->writePosition - tailBlock->data);

    public BlockBuffer(int allocationSize)
    {
        this = default;
        oneBlockSize = GetOneBlockSize(allocationSize);
        var newBlock = NewBlock(ref this);
    }

    private static int GetOneBlockSize(int x)
    {
        x = x > minAllocationSize ? x : minAllocationSize;
        return (x + 1024 - 1) / 1024 * 1024; // 向上取整为 1024 的倍数
    }

    public void Dispose()
    {
        var curBlock = headBlock;
        while (curBlock != null)
        {
            var next = curBlock->next;
            Marshal.FreeHGlobal((IntPtr)curBlock);
            curBlock = next;
        }

        headBlock = null;
        tailBlock = null;
        blockCount = 0;
    }

    public (ulong, ulong) GetBlockDataAddrRange()
    {
        if (headBlock == null) return (0, 0);
        return ((ulong)headBlock->data, (ulong)headBlock->writePosition);
    }


    public T* Malloc<T>() where T : unmanaged
    {
        int size = sizeof(T);
        T* obj = (T*)Allocate(size);
        *obj = default;
        return obj;
    }

    public void MallocBytes(byte[] bytes)
    {
        var ptr = Allocate(bytes.Length);
        fixed (byte* src = bytes)
        {
            Buffer.MemoryCopy(src, ptr, bytes.Length, bytes.Length);
        }
    }

    public void MallocNativeSimpleArray<T>(out NativeSimpleArray<T> outArray, int length, bool clearMemory = true) where T : unmanaged
    {
        if (length <= 0) throw new ArgumentOutOfRangeException(nameof(length));

        int size = sizeof(T) * length;
        var buffer = Allocate(size);

        if (clearMemory)
        {
            byte* p = buffer;
            for (int i = 0; i < size; i++) p[i] = 0;
        }

        outArray.m_Buffer = buffer;
        outArray.m_Length = length;
    }

    public void MallocNativeSimpleArray<T>(out NativeSimpleArray<T> outArray, T[] array) where T : unmanaged
    {
        if (array == null) throw new ArgumentNullException(nameof(array));

        MallocNativeSimpleArray(out outArray, array.Length, clearMemory: false);
        outArray.CopyFrom(array);
    }

    private byte* Allocate(int size)
    {
        size = (size + 7) & ~7; // 8字节对齐

        if (size > MaxDataSizeOfBlock)
            throw new OutOfMemoryException($"BlockBuffer: size {size} exceeds block capacity {MaxDataSizeOfBlock}");

        byte* ptr = tailBlock->writePosition;
        var endPtr = tailBlock->writePosition + size;

        if (endPtr > tailBlock->blockEnd)
        {
            var newBlock = NewBlock(ref this);
            ptr = newBlock->writePosition;
            newBlock->writePosition += size;
        }
        else
        {
            tailBlock->writePosition = endPtr;
        }

        return ptr;
    }

    private static Block* NewBlock(ref BlockBuffer buffer)
    {
        IntPtr mem = Marshal.AllocHGlobal(buffer.oneBlockSize);
        Block* block = (Block*)mem;

        block->blockEnd = (byte*)block + buffer.oneBlockSize - sizeof(byte*);
        block->writePosition = block->data;

        if (buffer.tailBlock == null)
        {
            buffer.headBlock = block;
            buffer.tailBlock = block;
            block->next = null;
            block->prev = null;
        }
        else
        {
            buffer.tailBlock->next = block;
            block->prev = buffer.tailBlock;
            block->next = null;
            buffer.tailBlock = block;
        }

        buffer.blockCount++;
        return block;
    }
}